from multiprocessing import get_context
from multiprocessing.managers import SyncManager
from typing import Callable, Optional

from infection_monkey.exploit import IAgentBinaryRepository
from infection_monkey.network import TCPPortSelector

from .http_agent_binary_server import AgentBinaryHTTPHandlerFactory, HTTPAgentBinaryServer


class HTTPAgentBinaryServerWithManagerFactory(HTTPAgentBinaryServer):
    """
    An HTTPAgentBinaryServer that has a SyncManager to generate events

    Note: Each instance of the server generated from this factory will have its own manager
          process, since the manager cannot be shared between processes.
    """

    def __init__(
        self,
        tcp_port_selector: TCPPortSelector,
        get_http_handler: AgentBinaryHTTPHandlerFactory,
    ):
        manager = get_context("spawn").Manager()
        create_event = manager.Event
        super().__init__(tcp_port_selector, get_http_handler, create_event, manager.Lock())


class HTTPAgentBinaryServerFactory:
    """
    Creates instances of HTTPAgentBinaryServer

    Each instance will run in the same managed process.
    """

    def __init__(
        self,
        tcp_port_selector: TCPPortSelector,
        agent_binary_repository: IAgentBinaryRepository,
        get_http_handler: Callable[[IAgentBinaryRepository], AgentBinaryHTTPHandlerFactory],
    ):
        self._tcp_port_selector = tcp_port_selector
        self._agent_binary_repository = agent_binary_repository
        self._get_http_handler = get_http_handler
        self._manager: Optional[SyncManager] = None

    def _get_manager(self) -> SyncManager:
        if self._manager is None:
            SyncManager.register("HTTPAgentBinaryServer", HTTPAgentBinaryServerWithManagerFactory)
            manager = get_context("spawn").Manager()
            self._manager = manager
            return manager

        return self._manager

    def __call__(self) -> HTTPAgentBinaryServer:
        manager = self._get_manager()
        return manager.HTTPAgentBinaryServer(  # type: ignore[attr-defined]
            self._tcp_port_selector,
            self._get_http_handler(self._agent_binary_repository),
        )
